"""
SQLite数据库管理模块
基于PRD精简版的数据库设计要求
"""

import sqlite3
import pandas as pd
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional, Dict, Any, List
import json
import logging

from config import DATABASE_CONFIG, CACHE_CONFIG

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self, db_path: str = None):
        """初始化数据库管理器"""
        self.db_path = db_path or DATABASE_CONFIG['db_path']
        self.init_database()
    
    def init_database(self):
        """初始化数据库表结构"""
        with sqlite3.connect(self.db_path) as conn:
            # 创建股票基本信息表
            conn.execute("""
                CREATE TABLE IF NOT EXISTS stocks (
                    stock_code TEXT PRIMARY KEY,
                    stock_name TEXT,
                    market TEXT,
                    industry TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # 创建资产负债表数据表
            conn.execute("""
                CREATE TABLE IF NOT EXISTS balance_sheet (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    stock_code TEXT NOT NULL,
                    report_date DATE NOT NULL,
                    report_type TEXT NOT NULL,  -- 季报/半年报/年报
                    account_name TEXT NOT NULL,  -- 科目名称
                    amount REAL NOT NULL,       -- 金额
                    category TEXT,              -- 分类标签
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (stock_code) REFERENCES stocks (stock_code),
                    UNIQUE(stock_code, report_date, account_name)
                )
            """)
            
            # 创建分析结果缓存表
            conn.execute("""
                CREATE TABLE IF NOT EXISTS analysis_cache (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    stock_code TEXT NOT NULL,
                    analysis_type TEXT NOT NULL,  -- 分析类型
                    analysis_data TEXT NOT NULL,  -- JSON格式的分析结果
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    expires_at TIMESTAMP NOT NULL,
                    FOREIGN KEY (stock_code) REFERENCES stocks (stock_code),
                    UNIQUE(stock_code, analysis_type)
                )
            """)
            
            # 创建索引提高查询性能
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_balance_sheet_stock_date 
                ON balance_sheet (stock_code, report_date, category)
            """)
            
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_analysis_cache_stock_type 
                ON analysis_cache (stock_code, analysis_type)
            """)
            
            conn.commit()
            logger.info("数据库初始化完成")
    
    def save_stock_info(self, stock_code: str, stock_name: str = None, 
                       market: str = None, industry: str = None):
        """保存股票基本信息"""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                INSERT OR REPLACE INTO stocks 
                (stock_code, stock_name, market, industry, updated_at)
                VALUES (?, ?, ?, ?, ?)
            """, (stock_code, stock_name, market, industry, datetime.now()))
            conn.commit()
    
    def save_balance_sheet_data(self, stock_code: str, data: pd.DataFrame):
        """保存资产负债表数据"""
        if data.empty:
            logger.warning(f"股票 {stock_code} 的数据为空，跳过保存")
            return
        
        with sqlite3.connect(self.db_path) as conn:
            # 删除旧数据
            conn.execute("""
                DELETE FROM balance_sheet WHERE stock_code = ?
            """, (stock_code,))
            
            # 插入新数据
            for _, row in data.iterrows():
                # 处理日期格式
                report_date = row.get('报告日期')
                if pd.notna(report_date):
                    if hasattr(report_date, 'strftime'):
                        report_date = report_date.strftime('%Y-%m-%d')
                    else:
                        report_date = str(report_date)
                else:
                    report_date = None

                conn.execute("""
                    INSERT INTO balance_sheet
                    (stock_code, report_date, report_type, account_name, amount, category)
                    VALUES (?, ?, ?, ?, ?, ?)
                """, (
                    stock_code,
                    report_date,
                    row.get('报告类型', '年报'),
                    row.get('科目'),
                    row.get('金额', 0),
                    row.get('category', 'unknown')
                ))
            
            conn.commit()
            logger.info(f"保存股票 {stock_code} 的资产负债表数据，共 {len(data)} 条记录")
    
    def get_balance_sheet_data(self, stock_code: str, 
                              start_date: str = None, 
                              end_date: str = None) -> pd.DataFrame:
        """获取资产负债表数据"""
        query = """
            SELECT stock_code, report_date, report_type, account_name, 
                   amount, category, created_at
            FROM balance_sheet 
            WHERE stock_code = ?
        """
        params = [stock_code]
        
        if start_date:
            query += " AND report_date >= ?"
            params.append(start_date)
        
        if end_date:
            query += " AND report_date <= ?"
            params.append(end_date)
        
        query += " ORDER BY report_date DESC, account_name"
        
        with sqlite3.connect(self.db_path) as conn:
            df = pd.read_sql_query(query, conn, params=params)
        
        return df
    
    def save_analysis_result(self, stock_code: str, analysis_type: str, 
                           analysis_data: Dict[str, Any], cache_days: int = 7):
        """保存分析结果到缓存"""
        expires_at = datetime.now() + timedelta(days=cache_days)
        
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                INSERT OR REPLACE INTO analysis_cache 
                (stock_code, analysis_type, analysis_data, created_at, expires_at)
                VALUES (?, ?, ?, ?, ?)
            """, (
                stock_code, 
                analysis_type, 
                json.dumps(analysis_data, ensure_ascii=False),
                datetime.now(),
                expires_at
            ))
            conn.commit()
    
    def get_analysis_result(self, stock_code: str, 
                          analysis_type: str) -> Optional[Dict[str, Any]]:
        """获取缓存的分析结果"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute("""
                SELECT analysis_data, expires_at 
                FROM analysis_cache 
                WHERE stock_code = ? AND analysis_type = ?
            """, (stock_code, analysis_type))
            
            result = cursor.fetchone()
            
            if result:
                analysis_data, expires_at = result
                expires_at = datetime.fromisoformat(expires_at)
                
                # 检查是否过期
                if datetime.now() < expires_at:
                    return json.loads(analysis_data)
                else:
                    # 删除过期数据
                    conn.execute("""
                        DELETE FROM analysis_cache 
                        WHERE stock_code = ? AND analysis_type = ?
                    """, (stock_code, analysis_type))
                    conn.commit()
        
        return None
    
    def cleanup_expired_cache(self):
        """清理过期的缓存数据"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute("""
                DELETE FROM analysis_cache 
                WHERE expires_at < ?
            """, (datetime.now(),))
            
            deleted_count = cursor.rowcount
            conn.commit()
            
            if deleted_count > 0:
                logger.info(f"清理了 {deleted_count} 条过期缓存记录")
    
    def get_database_stats(self) -> Dict[str, Any]:
        """获取数据库统计信息"""
        with sqlite3.connect(self.db_path) as conn:
            stats = {}
            
            # 股票数量
            cursor = conn.execute("SELECT COUNT(*) FROM stocks")
            stats['stocks_count'] = cursor.fetchone()[0]
            
            # 资产负债表记录数
            cursor = conn.execute("SELECT COUNT(*) FROM balance_sheet")
            stats['balance_sheet_records'] = cursor.fetchone()[0]
            
            # 缓存记录数
            cursor = conn.execute("SELECT COUNT(*) FROM analysis_cache")
            stats['cache_records'] = cursor.fetchone()[0]
            
            # 数据库文件大小
            if Path(self.db_path).exists():
                stats['db_size_mb'] = Path(self.db_path).stat().st_size / (1024 * 1024)
            else:
                stats['db_size_mb'] = 0
        
        return stats


# 全局数据库管理器实例
db_manager = DatabaseManager()
